"""
使用totalsegmentator 对processed中的数据进行分割
"""

import subprocess
import os
import glob
from tqdm import tqdm


def run_cli_command(inpath, outpath):
    command = [
        "TotalSegmentator",
        "-i",
        inpath,
        "-o",
        outpath,
        "-ta",
        "total_mr",
        "--device",
        "gpu:3",
    ]
    # 等的找找是哪个id,现在跑全身有些慢

    try:
        print("正在执行命令：", " ".join(command))
        result = subprocess.run(
            command, check=True, capture_output=True, text=True, env=os.environ.copy()
        )
        print("命令执行成功。输出如下：")
        print(result.stdout)
    except subprocess.CalledProcessError as e:
        print("命令执行失败。错误如下：")
        print(e.stderr)
        print("命令的标准输出：")
        print(e.stdout)
    except FileNotFoundError:
        print("命令未找到，请确保 'TotalSegmentator' 已安装并在系统 PATH 中。")
    except Exception as e:
        print(f"发生了一个未预料的错误：{e}")


if __name__ == "__main__":
    subject_id = "CHEN-REN-GENG_tra"
    source_dir = f"../prediction/{subject_id}/processed"
    files = glob.glob(os.path.join(source_dir, "**", "*.nii.gz"), recursive=True)
    print(len(files))  # ok
    save_dir = os.path.join("../segmentation", subject_id)
    if os.path.exists(save_dir):
        pass
    else:
        os.makedirs(save_dir, exist_ok=True)
        print("make dir ok")
    # 现在开始进行全分割了
    prefix = "seg"  # os.path.join(save_dir,)
    # print(files)
    for file in tqdm(files, desc="Processing files"):
        print("Start to segment the {}".format(file))
        classname = os.path.basename(file).split(".")[0]
        outpath = os.path.join(save_dir, prefix + "_" + classname)
        run_cli_command(file, outpath)
